package com.zx.learn.flink.source.custom;

import com.zx.learn.flink.jdbc.MySQLInfo;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.util.concurrent.TimeUnit;

/**
 * 实现flink与mysql的整合，读取mysql的数据并输出
 */
public class MysqlSource {
    public static void main(String[] args) throws Exception {
        // 1）获取flink流处理的运行环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        //env.setRuntimeMode(RuntimeExecutionMode.BATCH);
        System.out.println("执行环境默认的并行度："+env.getParallelism());
        // 2）接入自定义数据源
        //The parallelism of non parallel operator must be 1.
        //DataStreamSource<Order> streamSource = env.addSource(new MySource()).setParallelism(2);
        DataStreamSource<Student> streamSource = env.addSource(new MysqlSourceFunction()).setParallelism(1);
        System.out.println("streamSource的并行度："+streamSource.getParallelism());
        // 3）打印输出
        streamSource.print();

        // 4）启动运行作业
        env.execute();
    }

    @Data
    @AllArgsConstructor
    @NoArgsConstructor
    static class Student {
        private int id;
        private String name;
        private int age;
    }

    /**
     * 自定义数据源，读取mysql的数据
     */
    private static class MysqlSourceFunction extends RichParallelSourceFunction<Student> {
        //定义mysql的连接对象
        private Connection connection = null;
        //定义mysql的statement对象
        private PreparedStatement statement = null;
        //定义是否停止运行的标记

        /**
         * 实例化的时候会被执行一次，多个并行度会执行多次，因为有多个实例
         * 一般用于资源的初始化操作
         */
        @Override
        public void open(Configuration parameters) throws Exception {
            super.open(parameters);
            //注册驱动
            Class.forName("com.mysql.jdbc.Driver");

            //实例化connection的连接对象
            connection = DriverManager.getConnection(MySQLInfo.URL,MySQLInfo.USERNAME,MySQLInfo.PASSWORD);

            //实例化statement对象
            statement = connection.prepareStatement("select `id`,`name`,age from sample_user");
        }

        /**
         * 释放资源
         * @throws Exception
         */
        @Override
        public void close() throws Exception {
            super.close();
//            if(statement != null && !statement.isClosed()) statement.close();
//            if(connection != null && !connection.isClosed()) connection.close();
        }

        /**
         * 查询数据的核心方法
         * @param ctx
         * @throws Exception
         */
        @Override
        public void run(SourceContext<Student> ctx) throws Exception {
            while (true) {
                ResultSet resultSet = statement.executeQuery();

                while (resultSet.next()) {
                    int id = resultSet.getInt("id");
                    String name = resultSet.getString("name");
                    String passWord = resultSet.getString("password");
                    int age = resultSet.getInt("age");

                    //返回数据
                    ctx.collect(new Student(id, name, age));
                }
                resultSet.close();

                TimeUnit.SECONDS.sleep(1);
            }
        }

        @Override
        public void cancel() {

        }
    }
}
